# -*- coding: utf-8 -*-
# @Time   : 2019/11/19 下午1:58
# @Author : chujie
# @Email  : zjchenchujie@126.com
# @File   : generate_subset.py
# @Brief  : generate a subset from voc-formated dataset
"""
generate a subset from voc-formated dataset
"""

import os
import shutil
import argparse
import math


def select_subset(src_root_dir, dest_root_dir, fraction=0.3):
    xml_files = os.listdir(os.path.join(src_root_dir, "Annotations"))
    file_num = len(xml_files)
    step = int(1/fraction)
    dest_annotation_dir = os.path.join(dest_root_dir, "Annotations")
    dest_JPEGImages_dir = os.path.join(dest_root_dir, "JPEGImages")
    if os.path.exists(dest_annotation_dir):
        shutil.rmtree(dest_annotation_dir)
        os.makedirs(dest_annotation_dir)
    else:
        os.makedirs(dest_annotation_dir)

    if os.path.exists(dest_JPEGImages_dir):
        shutil.rmtree(dest_JPEGImages_dir)
        os.makedirs(dest_JPEGImages_dir)
    else:
        os.makedirs(dest_JPEGImages_dir)

    processed_num = 0
    for idx in range(1, file_num, step):
        file = xml_files[idx]
        file_name = file.split(".")[0]
        image_file = "{}.jpg".format(file_name)
        image_file_full_path = os.path.join(src_root_dir, "JPEGImages", image_file)
        xml_file_full_path = os.path.join(src_root_dir, "Annotations", file)
        dest_img_full_path = os.path.join(dest_JPEGImages_dir, image_file)
        shutil.copy(image_file_full_path, dest_img_full_path)
        shutil.copy(xml_file_full_path, dest_annotation_dir)
        processed_num += 1
        if processed_num % 100 == 0:
            print("Processed: {}, total: {}".format(processed_num, file_num*fraction))
    print("*"*80)
    print("Processed: {}, total: {}".format(processed_num, math.ceil(file_num * fraction)))


if __name__ == "__main__":
    PARSER = argparse.ArgumentParser(description="arguement")
    PARSER.add_argument("--src_root_dir", type=str, default="/home/chujie/edisk/dataset/VOCdevkit/VOC2012",
                        help="source voc dataset directory")
    PARSER.add_argument("--dest_root_dir", type=str, default="/home/chujie/PycharmProjects/data_tool/data/voc2012",
                        help="destination voc dataset directory")
    ARGS = PARSER.parse_args()

    select_subset(ARGS.src_root_dir, ARGS.dest_root_dir, 0.01)
